import random
import copy
import json
from random import randint
import rag
from FlagEmbedding import FlagModel

import LLM_prompt


class Dataset:
    def __init__(self, examples, negative_rate=2, positive_rate=1):
        # examples就是text（症状描述）和med（方剂的药物组成）两部分
        self.herb_vocab = self.build_herb_vocab()
        self.text = examples["text"]
        self.meds = examples["med"]
        self.samples = self.build_contrastive_samples(negative_rate, positive_rate, negative_sample_type="combine")

    @staticmethod
    def build_herb_vocab():
        herb_vocab = []
        for line in open("data/herb_vocab.txt"):
            tem = line.strip().split(' ')
            if len(tem) != 2:
                continue
            herb, count = tem
            if int(count) >= 5:
                herb_vocab.append(herb)
        return herb_vocab

    def write_samples(self):
        with open("data/contrastive_samples.jsonl", "w", encoding="utf8") as writer:
            for sample in self.samples:
                assert min([len(s) for s in sample]) > 0
                line_data = {"query": sample[0], "pos": sample[1],
                             "neg": sample[2]}
                writer.write(json.dumps(line_data, ensure_ascii=False) + "\n")

    @staticmethod
    def combine_positive_and_negative_sample(anchor_text, anchor_med, positive_meds, negative_meds, negative_texts):
        assert len(anchor_text) > 0

        def med_to_str(meds):
            assert len(meds) > 0
            return " ".join(meds)

        def med_list_to_str(med_list):
            assert len(med_list) > 0
            return [med_to_str(m) for m in med_list]

        results = []
        if len(negative_meds) > 0 and sum([len(s) for s in negative_meds]) > 0:
            results.append((anchor_text, [med_to_str(anchor_med)],
                            med_list_to_str(negative_meds)))
            for pos_med in positive_meds:
                results.append(
                    (anchor_text, [med_to_str(pos_med)], med_list_to_str(negative_meds)))
        if len(negative_texts) > 0:
            assert max([len(t) for t in negative_texts]) > 0
            results.append((med_to_str(anchor_med), [anchor_text], negative_texts))
            for pos_med in positive_meds:
                results.append((med_to_str(pos_med), [anchor_text], negative_texts))
        return results

    def build_contrastive_samples(self, positive_rate, negative_rate, negative_sample_type):
        # 以下构造负例
        results = []
        omitted_empty_num = 0  # 有些症状或者方剂是缺失的
        for i in range(len(self.text)):
            if len(self.meds[i]) == 0 or len(self.text[i]) == 0:
                omitted_empty_num += 1
                continue
            flag = 0
            text, meds = self.text[i], copy.deepcopy(self.meds[i])
            assert len(text) > 0 and len(meds) > 0
            positive_med_candidates = []
            negative_med_candidates = []
            negative_text_candidates = []
            for _ in range(min(positive_rate, len(meds))):
                if positive_rate > 1:
                    # 如果多次正采样，那么对方剂内药物顺序进行随机打乱
                    random.shuffle(meds)
                    positive_med_candidates.append(copy.deepcopy(meds))
            for _ in range(negative_rate):
                if negative_sample_type == "combine":
                    # 两种策略各半
                    if random.random() < 0.3:
                        flag = 1
                    elif random.random() > 0.7:
                        flag = 0
                    else:
                        flag = 2
                if negative_sample_type == "random" or len(self.meds[i]) <= 2 or len(text) <= 3 or flag == 1:
                    # 随机整体替换,区分替换药物还是替换症状是因为有些项的药物或者症状是缺少的
                    _, fake_med = self.negative_sample_prescription(i, "med")
                    negative_med_candidates.append(fake_med)
                elif negative_sample_type == "replace" or flag == 2:
                    # 替换部分药物
                    med = self.meds[i]
                    fake_med, to_replace_id = self.replace_random_med(copy.deepcopy(med), self.herb_vocab)
                    negative_med_candidates.append(fake_med)
                else:
                    fake_text, med = self.negative_sample_prescription(i, "text")
                    assert len(fake_text) > 0
                    negative_text_candidates.append(fake_text)
            results.extend(
                self.combine_positive_and_negative_sample(text, meds, positive_med_candidates, negative_med_candidates,
                                                          negative_text_candidates))
        return results

    def negative_sample_prescription(self, real_index, sample_type):
        if sample_type == "med":
            fake_med_id = random.choice(range(len(self.text)))
            while fake_med_id == real_index or len(self.meds[fake_med_id]) == 0:
                fake_med_id = random.choice(range(len(self.text)))
            text, fake_med = self.text[real_index], self.meds[fake_med_id]
            real_med = self.meds[real_index]
            assert len(fake_med) > 0
        else:
            fake_text_id = random.choice(range(len(self.text)))
            while fake_text_id == real_index or len(self.text[fake_text_id]) == 0:
                fake_text_id = random.choice(range(len(self.text)))
            text, fake_med = self.text[fake_text_id], self.meds[real_index]
            real_med = self.meds[fake_text_id]
            assert len(text) > 0

        return text, fake_med

    @staticmethod
    def replace_random_med(med, herb_vocab):
        if len(med) <= 4:
            rep_med = random.choice(herb_vocab)
            to_replace_id = random.choice(range(len(med)))
            med[to_replace_id] = rep_med
            to_replace_id = [to_replace_id]
        else:
            replace_num = random.choice(range(2, round(len(med) / 2) + 1))
            to_replace_id = random.sample(range(len(med)), replace_num)
            rep_meds = random.sample(herb_vocab, replace_num)
            for i, m in zip(to_replace_id, rep_meds):
                med[i] = m
        return med, to_replace_id

    @staticmethod
    def find_difference(med, fake):
        to_replace_id = []
        med_set = set(med)
        for i, m in enumerate(fake):
            if m not in med_set:
                to_replace_id.append(i)
        return to_replace_id

    @staticmethod
    def get_similar_negative_samples(text, med, compare_text, compare_meds, sample_num=20):
        candidates = []
        duplicate_num = 0
        for i in range(len(compare_text)):
            if len(compare_meds[i]) == 0:
                continue
            js = Dataset.jaccard(text, compare_text[i])
            if js == 1:
                duplicate_num += 1
            if js == 0 or js == 1 or med == compare_meds[i]:
                continue
            candidates.append((js, compare_text[i], compare_meds[i]))
            candidates.sort(key=lambda k: k[0], reverse=True)
            candidates = candidates[:sample_num]
        return candidates, duplicate_num

    @staticmethod
    def jaccard(t1, t2):
        list1, list2 = set(list(t1)), set(list(t2))
        intersection = list1.intersection(list2)
        union = list1.union(list2)
        return len(intersection) / float(len(union))


def read_prescription_pair_data(fname):
    data = {"text": [], "med": []}
    for line in open(fname):
        sample = json.loads(line.strip())
        data["text"].append(sample["text"])
        data["med"].append(sample["med"])
    return data


class DosageDataset:
    """
    验证reranker的质量可以考虑对古代剂量和现代剂量分别计算，古代剂量数据较多，可以拿出更多数据作为测试，现代剂量可以使用相对少的数据进行测试
    可以考虑使用大模型进行评价并比较和大模型的一致性（古代和现代）
    """

    def __init__(self, records, bge_model=None):
        self.records = records  # 主治，组成， 药物， 剂量
        self.herb_vocab = Dataset.build_herb_vocab()
        self.bge_model = bge_model
        if bge_model is not None:
            self.faiss_rag = self.build_faiss_base(bge_model)
        else:
            self.faiss_rag = None

    def build_faiss_base(self, bge_model):
        faiss_rag = rag.FaissRAG(self.records, bge_model, 1024)
        return faiss_rag

    @staticmethod
    def swap(med_dose: dict):
        # swap the doses of two herbs with different dosages
        ids = list(range(len(med_dose)))
        random.shuffle(ids)
        flag = False
        for i in ids:
            for j in range(len(med_dose)):
                if med_dose[i] != med_dose[j]:
                    med_dose[i], med_dose[j] = med_dose[j], med_dose[i]
                    flag = True
                    break
            if flag:
                break
        return med_dose

    @staticmethod
    def change_num(med_dose: dict):
        """
        目前先只考虑现代剂量
        """
        change_herb = random.choice(med_dose.keys())
        original_dose, original_unit = med_dose[change_herb]
        original_dose = float(original_dose)
        max_dose = max([float(dose[0]) for dose in med_dose.values() if dose is not None])
        if original_dose is None:
            dose = randint(1, int(max_dose))
            while dose == original_dose:
                dose = randint(1, int(max_dose))
        else:
            dose = randint(1, int(original_dose * 2))
            while dose == original_dose:
                dose = randint(1, int(original_dose * 2))
        med_dose[change_herb] = dose, original_unit
        return med_dose

    @staticmethod
    def multiply(med_dose: dict):
        multiplier_flag = True
        doses = med_dose.values()
        for dose in doses:
            if dose[0] % 2 != 0:
                multiplier_flag = False
                break
        if multiplier_flag:
            med_dose = {med: (med_dose[med][0] / 2, med_dose[med][1]) for med in med_dose}
        else:
            med_dose = {med: (med_dose[med] * 2, med_dose[med][1]) for med in med_dose}
        return med_dose

    @staticmethod
    def combine_med_and_dose(record):
        return {med: dose for med, dose in zip(record["药物"], record[u"剂量"])}

    def build_contrastive_samples(self):
        # 负例：随机调换药物剂量，随机增加或减少药物剂量，负例也可以由大模型预测得到
        # 正例：药物剂量乘以系数（尽量保持整数），
        # todo: 大模型预测做负例
        contrastive_pairs = []
        for rid, record in enumerate(self.records):
            cure_text = record[u"主治"]
            constituent = record[u"组成"]
            anchor = self.combine_med_and_dose(record)
            med_doses = {med: LLM_prompt.extract_num_and_unit(anchor[med]) for med in anchor}

            pos_1 = self.multiply(med_doses)
            # bge retrive pos and neg
            neg_1 = self.swap(med_doses)
            contrastive_pairs.append((cure_text, anchor, pos_1, neg_1))
            if self.faiss_rag is not None:
                cure_vec = self.bge_model.encode(cure_text)
                cure_ids, cure_records = self.faiss_rag.retrieve(cure_vec, topk=10, retrieve_strategy="cure")
                constituent_vec = self.bge_model.encode(constituent)
                med_ids, med_records = self.faiss_rag.retrieve(constituent_vec, topk=10,
                                                               retrieve_strategy="constituent")
                assert cure_ids[1] != rid
                contrastive_pairs.append((cure_text, anchor, self.combine_med_and_dose(self.records[cure_ids[1]]),
                                          self.combine_med_and_dose(self.records[cure_ids[-1]])))
                contrastive_pairs.append((cure_text, anchor, self.combine_med_and_dose(self.records[med_ids[1]]),
                                          self.combine_med_and_dose(self.records[med_ids[-1]])))
        return contrastive_pairs


def build_contrastive_data_for_bge():
    data = read_prescription_pair_data("data/prescription_pairs_text_total.json")
    dataset = Dataset(data)
    dataset.write_samples()


def build_contrastive_data_for_dose(records, bge_path="../bge-large-zh-v1.5/tcm_dosage"):
    bge_model = FlagModel(bge_path,
                          use_fp16=True)  # Setting use_fp16 to True speeds up computation with a slight performance degradation
    dataset = DosageDataset(records, bge_model)
    pairs = dataset.build_contrastive_samples()     # (anchor, pos, neg)
    return pairs


if __name__ == "__main__":
    pass